#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time     : 2024/7/5 13:50
# @Author   : 22759
# @Email    : zjh13092921608@163.com
# @Project  : S2
# @File     : batchTask.py
# @Software : PyCharm
import os
import base64
import traceback
from concurrent.futures import ThreadPoolExecutor

from utils.cryption import Cryption
from celery_tasks.celeryUtils import SQLServerClass, RemoteRequest


class BatchCenter(object):
    """批量中枢"""

    def __init__(self, _type, _ids, _bhr, _batch_id):
        self.type, self.ids, self.behavior, self.batch_id = _type, _ids, _bhr, _batch_id
        self.url, self.path, self.file, self.statement = '', '/storage/static/batch/', self._filename(), ''
        self.sql_server = SQLServerClass(
            host='localhost',
            port=3306,
            username='insec',
            password='InSec888!',
            database='safeplatform',
        )

    def __del__(self):
        self.sql_server.close()

    def center(self):
        print('Task action')
        if self.behavior == 1:
            self._up_system()
        elif self.behavior == 2:
            self._up_library()
        elif self.behavior == 3:
            self._reboot()
        else:
            print('Task is not OK')
            return False
        print('Task is OK')
        return True

    def _filename(self):
        """文件名函数"""
        if self.behavior == 1:
            if 'T3' in self.type:
                return 'upgrade_system.bin'
            elif any(mark in self.type for mark in {'A4', 'LAS'}):
                return 'upgrade_system-update.tar.gz'
        elif self.behavior == 2 and any(mark in self.type for mark in {'T3', 'A4', 'LAS'}):
            return 'upgrade_library.gpg'
        return False

    def _up_system(self):
        """升级系统"""
        if not self.file:
            return False
        if 'A4' in self.type:
            """流量审计系统升级"""
            self.url = 'https://{}/api/api_v1/sys_upgrade/remote_update/'
        elif 'T3' in self.type:
            """防火墙系统升级"""
            self.url = 'https://{}/cgi-bin/webui/rpc/upgrade'
        elif 'LAS' in self.type:
            """日志审计系统升级"""
            self.url = 'https://{}/api/licenseInformation/shard_upgrade'
        for index, _id in enumerate(self.batch_id):
            """升级系统"""
            with ThreadPoolExecutor(max_workers=4) as executor:
                executor.submit(self._update_device, _id, index)
        """升级系统任务执行完成，删除升级文件，调试阶段暂不删除文件"""
        try:
            os.remove(os.path.join(self.path, self.file))
        except FileNotFoundError:
            print('文件不存在')
        except Exception as error:
            print(error)

    def _up_library(self):
        """升级特征库"""
        if not self.file:
            return False
        if 'A4' in self.type:
            """流量审计特征库升级"""
            self.url = 'https://{}/api/api_v1/sys_upgrade/remote_library/'
        elif 'T3' in self.type:
            """防火墙特征库升级"""
            self.url = 'https://{}/cgi-bin/webui/rpc/upgrade/'
        # elif 'LAS' in self.type:
        #     """日志审计特征库升级"""
        #     self.url = 'https://{}/api/licenseInformation/upgrade'
        for index, _id in enumerate(self.batch_id):
            """升级特征库"""
            with ThreadPoolExecutor(max_workers=4) as executor:
                executor.submit(self._update_device, _id, index)
        """升级特征库任务执行完成，删除升级文件，调试阶段暂不删除文件"""
        try:
            os.remove(os.path.join(self.path, self.file))
        except FileNotFoundError:
            print('文件不存在')
            return False
        except Exception as error:
            print(error)
            return False

    def _reboot(self):
        """重启"""

    def _device_info(self, _id):
        self.statement = f"""
        SELECT
            ip, username, passwd, isauthenticate
        FROM
            t_device
        WHERE
            id={_id};
        """
        self.sql_server.execute(self.statement)
        device_info = self.sql_server.fetchone()
        return device_info

    def _batch_info(self, _id, _status, _cause=''):
        self.statement = f"""
        UPDATE
            batch_operation
        SET
            status = '{_status}', cause = '{_cause}'
        WHERE 
            id = {_id};"""
        try:
            self.sql_server.execute(self.statement)
            return True
        except Exception as error:
            print(error, traceback.format_exc())
            return False

    def _update_device(self, batch_id, device_index):
        """设备升级"""
        status, message, password = -1, '升级失败，使用暂不支持升级的设备类型', None
        if self.url:
            """启动任务"""
            self._batch_info(batch_id, 0)
            """设备信息"""
            device_info = self._device_info(self.ids[device_index])
            """接口请求必要参数"""
            if device_info['passwd']:
                password = bytes.decode(Cryption.decrypt(device_info['passwd']), "utf-8").split(',')[0]
            if 'A4' in self.type:
                password = base64.b64encode(password.encode('utf-8'))
            data = self._data(device_info['username'], password)
            """调用升级"""
            try:
                request = RemoteRequest(_type=self.type, _ip=device_info['ip'], _behavior=self.behavior, _data=data)
                result = request.base_requests(
                    url=self.url.format(device_info['ip']),
                    method='post',
                    path=self.path,
                    filename=self.file,
                )
                """成功失败记录"""
                status = 1 if result['mark'] else -1
                message = result['message']
            except Exception as error:
                print(error, traceback.format_exc())
                status, message = -1, '升级失败，指令下发失败'
        self._batch_info(batch_id, status, message)

    def _data(self, username, password):
        """返回请求数据"""
        data = dict()
        if any(t in self.type for t in ('A4', 'LAS')):
            data = {'username': username, 'password': password}
        elif 'T3' in self.type:
            data = {'method': 'login', 'params': [username, password]}
        return data
